# tools.py

import psycopg2
import psycopg2.extras
import mysql.connector
import sqlite3
import oracledb  # 以前的 cx_Oracle

from config_app.models import DataSource, DataSourceConfig, Agent, SQLConfig

def execute_sql(sql_config_name: str, params: dict) -> dict:
    """
    支持多数据库的安全工具函数，用于执行预定义的SQL配置。
    它只接受SQL配置名称和参数，从不接受原始SQL。

    :param sql_config_name: SQLConfig的名称，例如 'query_user_with_orders'
    :param params: 一个包含所需参数的字典，例如 {'user_id': '123'}
    :return: 包含查询结果或错误的字典
    """
    try:
        sql_config = SQLConfig.objects.get(name=sql_config_name)
    except SQLConfig.DoesNotExist:
        return {"error": f"未找到名为'{sql_config_name}'的SQL配置。"}

    datasource = sql_config.datasource_id

    source = DataSource.objects.get(id=datasource)

    relevant_tables = get_relevant_tables(sql_config.sql)

    # print(relevant_tables)

    # 相关表格和字段的信息
    relevant_table_fields = []
    
    # 安全检查：确保调用者提供了所有必需的参数
    # print(f"Required params: {sql_config.params}")
    # print(f"Provided params: {params}")
    for req_param in sql_config.params.split(';'):
        if req_param not in params:
            return {"error": f"执行'{sql_config_name}'时缺少必要参数: {req_param}"}

    conn = None
    try:
        db_type = source.engine
        sql_template = sql_config.sql
        
        # --- 动态数据库连接和执行 ---

        if db_type == 'postgresql':
            conn = psycopg2.connect(
                host=source.host,
                port=source.port,
                user=source.user,
                password=source.password,
                dbname=source.database
            )
            
            # psycopg2 的 DictCursor 可以直接返回类字典对象
            with conn.cursor(cursor_factory=psycopg2.extras.DictCursor) as cursor:
                for table in relevant_tables:
                    cursor.execute("SELECT column_name, data_type, pg_catalog.col_description(c.oid, cols.ordinal_position) as column_comment FROM information_schema.columns cols JOIN pg_catalog.pg_class c ON c.relname = cols.table_name WHERE cols.table_name = %s AND cols.table_schema = 'public'", [table])
                    fields = [{
                        'name': row['column_name'],
                        'type': row['data_type'],
                        'description': row['column_comment'] or ''
                    } for row in cursor.fetchall()]
                    relevant_table_fields.append({'name': table, 'fields': fields})
                cursor.execute(sql_template, params)
                records = [dict(row) for row in cursor.fetchall()]

        elif db_type == 'mysql':
            conn = mysql.connector.connect(
                host=source.host,
                port=source.port,
                user=source.user,
                password=source.password,
                database=source.database
            )
            # mysql-connector-python 可以配置直接返回字典
            with conn.cursor(dictionary=True) as cursor:
                for table in relevant_tables:
                    # print([source.database, table])
                    cursor.execute("SELECT column_name, data_type, column_comment FROM information_schema.columns WHERE table_schema = %s AND table_name = %s", [source.database, table])
                    # print(cursor.fetchall())
                    # all_fields = cursor.fetchall()
                    fields = [{
                        'name': row['COLUMN_NAME'],
                        'type': row['DATA_TYPE'],
                        'description': row['COLUMN_COMMENT'] or ''
                    } for row in cursor.fetchall()]
                    # print(fields)
                    relevant_table_fields.append({'name': table, 'fields': fields})
                cursor.execute(sql_template, params)
                records = cursor.fetchall()

                from decimal import Decimal
                from datetime import datetime
                table_data = []
                # print(records)
                for row in records:
                    row_dict = {}
                    for col, val in row.items():
                        if isinstance(val, Decimal):
                            # 处理 Decimal 类型（如数据库中的 Decimal 字段）
                            row_dict[col] = float(val)
                        elif isinstance(val, datetime):
                            # 处理 datetime 类型，转为字符串
                            row_dict[col] = val.strftime('%Y-%m-%d %H:%M:%S')
                        else:
                            # 其他类型直接保留（如整数、None、字符串等）
                            row_dict[col] = val
                    table_data.append(row_dict)
                records = table_data
        
        elif db_type == 'sqlite':
            # SQLite 的参数风格是 :name 而不是 %(name)s，需要转换
            sql_for_sqlite = sql_template.replace('%(', ':').replace(')s', '')
            conn = sqlite3.connect(database=source.database) # dbname 存储文件路径
            conn.row_factory = sqlite3.Row  # 让返回的行可以像字典一样访问
            with conn.cursor() as cursor:
                cursor.execute(sql_for_sqlite, params)
                records = [dict(row) for row in cursor.fetchall()]

        elif db_type == 'oracle':
            # Oracle 的参数风格也是 :name
            sql_for_oracle = sql_template.replace('%(', ':').replace(')s', '')
            # Oracle 的 DSN (Data Source Name) 可能需要特殊构造
            dsn = oracledb.makedsn(source.host, source.port, service_name=source.database)
            conn = oracledb.connect(
                user=source.user,
                password=source.password,
                dsn=dsn
            )
            with conn.cursor() as cursor:
                cursor.execute(sql_for_oracle, params)
                # Oracle驱动不直接返回字典，需要手动构造
                columns = [col[0].lower() for col in cursor.description]
                records = [dict(zip(columns, row)) for row in cursor.fetchall()]
        
        else:
            return {"error": f"不支持的数据库类型: {db_type}"}

        # 脱敏逻辑可以在这里或之后添加
        # ...

        return {"success": True,"描述":relevant_table_fields,"data": records}

    except Exception as e:
        # 捕获所有可能的数据库异常
        return {"error": f"数据库查询执行失败 ({source.name}): {str(e)}"}
    
    finally:
        # 确保数据库连接总是被关闭
        if conn:
            conn.close()


def get_relevant_tables(sql):
    """
    从SQL语句中提取相关的表名。
    这里使用一个简单的正则表达式来匹配 FROM 子句后的表名。
    注意：这是一个基本实现，可能需要根据实际SQL语法进行调整。
    """
    import re
    # 匹配 FROM 子句后的表名 和 JOIN 子句后的表名
    pattern = r'\bFROM\s+(\w+)|JOIN\s+(\w+)'
    # 查找所有匹配项
    matches = re.findall(pattern, sql, re.IGNORECASE)

    # 返回一个表格的list
    tables = [match[0] or match[1] for match in matches]
    return tables

